# encoding: utf-8
"""
@author:  shaoqijun
@date:  2020/7/30
"""
from __future__ import print_function, division

import numpy as np
import torchvision
import matplotlib.pyplot as plt
import torch
from config import  cfg

#可视化一些图像
def imshow(inp, title=None):
    plt.ion()
    """Imshow for Tensor."""
    inp = inp.numpy().transpose((1, 2, 0))
    mean = np.array(cfg.TRAIN.MEAN)
    std = np.array(cfg.TRAIN.STD)
    inp = std * inp + mean
    inp = np.clip(inp, 0, 1)
    plt.imshow(inp)
    if title is not None:
        plt.title(title)
    plt.pause(1)  # pause a bit so that plots are updated
    plt.ioff()
    plt.show()

def show_img(dataloaders, class_names):
    # Get a batch of training data
    inputs, classes = next(iter(dataloaders['train']))

    # Make a grid from batch
    out = torchvision.utils.make_grid(inputs)


    #imshow(out, title=[class_names[x] for x in classes])
    imshow(out)

def sw_add_image(train_loader):
    inputs, classes = next(iter(train_loader))
    out = torchvision.utils.make_grid(inputs)
    # inp = out.numpy().transpose((1, 2, 0))
    # mean = np.array(cfg.TRAIN.MEAN)
    # std = np.array(cfg.TRAIN.STD)
    # inp = std * inp + mean
    # inp = np.clip(inp, 0, 1)
    # inp = inp.transpose((2,0,1))
    # output = torch.from_numpy(inp)
    #return output
    return inputs, out